/*
    This code configures a Windows process to only load Microsoft-signed DLLs,
    enhancing security by preventing the execution of unsigned or malicious code. 

    @5mukx...

    Resource Used: 
        * https://0x00sec.org/t/building-a-basic-dll-blocker/36049
        * https://www.ired.team/offensive-security/defense-evasion/preventing-3rd-party-dlls-from-injecting-into-your-processes
        * https://techcommunity.microsoft.com/t5/microsoft-intune/microsoft-recommended-block-rules-for-dlls/m-p/3670536 [community Platform]

*/

use std::{ffi::OsString, iter::once, os::windows::ffi::OsStrExt, path::Path, ptr::null_mut};
use winapi::{
    ctypes::c_void, 
    shared::minwindef::{DWORD, LPVOID}, 
    um::{
        errhandlingapi::GetLastError, 
        handleapi::CloseHandle, 
        heapapi::{GetProcessHeap, HeapAlloc, HeapFree}, 
        processthreadsapi::{CreateProcessA, CreateProcessW, DeleteProcThreadAttributeList, InitializeProcThreadAttributeList, UpdateProcThreadAttribute, PROCESS_INFORMATION, PROC_THREAD_ATTRIBUTE_LIST, STARTUPINFOW}, 
        winbase::{EXTENDED_STARTUPINFO_PRESENT, STARTUPINFOEXA, STARTUPINFOEXW},
    }
};

const PROC_THREAD_ATTRIBUTE_MITIGATION_POLICY: usize = 0x00020007;
const PROCESS_CREATION_MITIGATION_POLICY_BLOCK_NON_MICROSOFT_BINARIES_ALWAYS_ON: u64 = 0x00000001u64 << 44;

#[repr(transparent)]
#[allow(non_camel_case_types)]
pub struct STARTUPINFOW_FLAGS(pub u32);


#[allow(non_snake_case)]
unsafe fn CreateProcess_Dll_Blocker(){
    // let windir = std::env::var("WINDIR").unwrap();
    // let exe_path = format!("{}\\System32\\SystemSettingsBroker.exe", windir);
    // let wide_windir: Vec<u16> = OsString::from(exe_path).encode_wide().chain(once(0)).collect();

    // let exe_path = Path::new("C:\\Windows\\System32\\notepad.exe");
    // let wide_windir: Vec<u16> = exe_path.as_os_str().encode_wide().chain(std::iter::once(0)).collect();


    let exe_path = OsString::from("C:\\Windows\\System32\\SystemSettingsBroker.exe");
    let wide_exe_path: Vec<u16> = exe_path.encode_wide().chain(std::iter::once(0)).collect();

    // let mut si = STARTUPINFOEXW {
    //     StartupInfo: STARTUPINFOW {
    //         cb: std::mem::size_of::<STARTUPINFOW>() as DWORD,
    //         lpReserved: null_mut(),
    //         lpDesktop: null_mut(),
    //         lpTitle: null_mut(),
    //         dwX: 0,
    //         dwY: 0,
    //         dwXSize: 0,
    //         dwYSize: 0,
    //         dwXCountChars: 0,
    //         dwYCountChars: 0,
    //         dwFillAttribute: 0,
    //         dwFlags: EXTENDED_STARTUPINFO_PRESENT,
    //         wShowWindow: 0,
    //         cbReserved2: 0,
    //         lpReserved2: null_mut(),
    //         hStdInput: null_mut(),
    //         hStdOutput: null_mut(),
    //         hStdError: null_mut(),
    //     },
    //     lpAttributeList: null_mut(),
    // };

    let mut si: STARTUPINFOEXA = std::mem::zeroed();
    si.StartupInfo.cb = std::mem::size_of::<STARTUPINFOEXA>() as u32;
    si.StartupInfo.dwFlags = EXTENDED_STARTUPINFO_PRESENT;

    // let mut si: STARTUPINFOW = std::mem::zeroed();
    // si.cb = std::mem::size_of::<STARTUPINFOW>() as DWORD;

    let mut pi = PROCESS_INFORMATION {
        hProcess: null_mut(),
        hThread: null_mut(),
        dwProcessId: 0,
        dwThreadId: 0,
    }; 

    let mut attr_size = 0;
    InitializeProcThreadAttributeList(
        null_mut(),
        1,
        0,
        &mut attr_size,
    );

    let attr_list: *mut c_void = HeapAlloc(
        GetProcessHeap(), 
        8,
        attr_size
    );

    if attr_list.is_null() {
        println!("[-] HeapAlloc failed");
        return;
    }

    if InitializeProcThreadAttributeList(
        attr_list as *mut PROC_THREAD_ATTRIBUTE_LIST,
        1,
        0,
        &mut attr_size
    ) == 0 {
        println!("[-] InitializeProcThreadAttributeList failed: {}", GetLastError());
        HeapFree(GetProcessHeap(), 0, attr_list);
        return;
    }

    let policy = PROCESS_CREATION_MITIGATION_POLICY_BLOCK_NON_MICROSOFT_BINARIES_ALWAYS_ON;

    if UpdateProcThreadAttribute(
        attr_list as *mut PROC_THREAD_ATTRIBUTE_LIST,
        0,
        PROC_THREAD_ATTRIBUTE_MITIGATION_POLICY,
        &policy as *const _ as LPVOID,
        std::mem::size_of::<u64>() as usize,
        null_mut(),
        null_mut(),
    ) == 0 {
        println!("[-] UpdateProcThreadAttribute failed: {}", GetLastError());
        DeleteProcThreadAttributeList(attr_list as *mut PROC_THREAD_ATTRIBUTE_LIST);
        HeapFree(GetProcessHeap(), 0, attr_list);
        return;
    }

    si.lpAttributeList = attr_list as *mut PROC_THREAD_ATTRIBUTE_LIST;

    let result = CreateProcessA(
        null_mut(),
        wide_exe_path.as_ptr() as *mut i8,
        null_mut(),
        null_mut(),
        0,
        EXTENDED_STARTUPINFO_PRESENT,
        null_mut(),
        null_mut(),
        &mut si.StartupInfo,
        &mut pi
    );

    // let result = CreateProcessW(
    //     null_mut(),
    //     wide_exe_windir.as_ptr() as *mut u16,
    //     null_mut(),
    //     null_mut(),
    //     0,
    //     EXTENDED_STARTUPINFO_PRESENT,
    //     null_mut(),
    //     null_mut(),
    //     &mut si.StartupInfo,
    //     &mut pi
    // );

    if result == 0 {
        println!("[-] CreateProcess Failed: {}", GetLastError());
    } else {
        println!("[+] CreateProcess succeeded");
    }

    DeleteProcThreadAttributeList(attr_list as *mut PROC_THREAD_ATTRIBUTE_LIST);
    HeapFree(GetProcessHeap(), 0, attr_list);
    CloseHandle(pi.hProcess);
    CloseHandle(pi.hThread);
}

fn main(){
    unsafe {
        CreateProcess_Dll_Blocker();
    }
}

